import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(project_root)
import threading

#import open3d as o3d
import numpy as np
import numpy as np
from sklearn.neighbors import KDTree

from concurrent.futures import ThreadPoolExecutor

from utils import *

lat0 = None
lon0 = None
h0 = None

all_data = {'x': [], 'y': [], 'z': [], 's': [], 'dn': []}
all_data_lock = threading.Lock()

def process_mat_file(mat_data, x_range, y_range):
    global lat0, lon0, h0
    global all_data

    start_frame_llh = mat_data.get('StartFrameLLH', None)
    start_frame_prh = mat_data.get('StartFrameRPY', None)
    azi_map = mat_data.get('aziLinearlityMap', None)
    ele_map = mat_data.get('eleLinearlityMap', None)
    errorMap = mat_data.get('errorMap', None)#np.ones_like(ele_map) # 
    target_x = mat_data.get('targetX', None)
    target_y = mat_data.get('targetY', None)
    target_z = mat_data.get('targetZ', None)
    dn = mat_data.get('DN', None)

    if not all([azi_map is not None, ele_map is not None, errorMap is not None, target_x is not None, 
                target_y is not None, target_z is not None, dn is not None, start_frame_llh is not None, start_frame_prh is not None]):
        raise ".mat文件中缺少数据"
    
    lat = start_frame_llh[0][0]  # 纬度
    lon = start_frame_llh[0][1]  # 精度
    h = start_frame_llh[0][2] / 1000  # 海拔高度

    pitch = start_frame_prh[0][0]  # 俯仰角（重力基准）
    roll = start_frame_prh[0][1]  # 滚转角（重力基准）
    heading = start_frame_prh[0][2]  # 与北方向夹角

    if lat0 == None:
        lat0 = lat
        lon0 = lon
        h0 = h

    X = target_x.flatten()
    Y = target_y.flatten()
    Z = target_z.flatten()
    DN = dn.flatten()
    azi = azi_map.flatten()
    ele = ele_map.flatten()
    err = errorMap.flatten()

    s = azi * ele * err

    raw_data = {
        'x': X,
        'y': Y,
        'z': Z,
        'dn': DN,
        's': s,
        'aziLinearlityMap': azi,
        'eleLinearlityMap': ele,
        'errorMap': err,
    }
    
    filtered_data = filter_z(raw_data)
    filtered_data = filter_data(filtered_data, x_range, y_range, filter_noise=True)

    X = filtered_data['x']
    Y = filtered_data['y']
    Z = filtered_data['z']
    DN = filtered_data['dn']
    s = filtered_data['s']

    X, Y, Z = local_to_global_coords(X, Y, Z, lat, lon, h, pitch, roll, heading, lat0, lon0, h0)
    
    with all_data_lock:
        all_data['x'].extend(X)
        all_data['y'].extend(Y)
        all_data['z'].extend(Z)
        all_data['s'].extend(s)
        all_data['dn'].extend(DN)

def plot_3d():
    global lat0, lon0, h0, all_data

    # 提取点云数据
    X = np.array(all_data['x'])
    Y = np.array(all_data['y'])
    Z = np.array(all_data['z'])
    S = np.array(all_data['s'])
    DN = np.array(all_data['dn'])

    # 计算RGB颜色，假设compute_rgb返回的是[0, 255]范围内的RGB值
    rgb = compute_rgb(Z, S, DN)  # 应返回形状为(N, 3)的数组

    # 检查rgb的形状是否正确
    if rgb.shape[1] != 3:
        raise ValueError("compute_rgb函数应返回形状为(N, 3)的RGB数组")

    # 创建Open3D点云对象
    points = np.vstack((X, Y, Z)).T  # 形状为(N, 3)

    
    # 使用KDTree查找邻居
    tree = KDTree(points)

    # 定义距离阈值和最小邻居数
    distance_threshold = 0.1  # 邻居的最大距离
    min_neighbors = 2  # 最小邻居数量

    # 存储非孤立点的索引
    non_isolated_points = []
    total_points = len(points)
    progress_interval = 100

    for i in range(total_points):
        # 查找距离当前点在距离阈值内的邻居
        indices = tree.query_radius([points[i]], r=distance_threshold)[0]
        if len(indices) > min_neighbors:
            non_isolated_points.append(i)

        # 每100个点打印一次进度
        if (i + 1) % progress_interval == 0 or i == total_points - 1:
            print(f"Processed {i + 1}/{total_points} points...")
    # 保留非孤立点的坐标
    filtered_points = points[non_isolated_points]
    filtered_rgb = rgb[non_isolated_points]

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(filtered_points)

    # 设置点的颜色，归一化到[0, 1]
    pcd.colors = o3d.utility.Vector3dVector(filtered_rgb)

    # 创建坐标轴（长度可以根据点云范围调整）
    axis_length = max(np.ptp(X), np.ptp(Y), np.ptp(Z)) * 0.1  # 点云范围的10%
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_length)

    # 可视化设置
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name='3D 点云可视化', width=1280, height=720)
    vis.add_geometry(pcd)
    vis.add_geometry(coordinate_frame)

    # 设置背景颜色（可选）
    opt = vis.get_render_option()
    opt.background_color = np.asarray([0.1, 0.1, 0.1])  # 深色背景
    opt.point_size = 1  # 点的大小

    # 设置视图参数（可选）
    ctr = vis.get_view_control()
    ctr.set_zoom(0.8)
    ctr.set_front([0.0, 0.0, -1.0])  # 视角方向
    ctr.set_lookat([np.mean(X), np.mean(Y), np.mean(Z)])  # 观察点
    ctr.set_up([0.0, -1.0, 0.0])  # 上方向

    # 运行可视化
    vis.run()
    vis.destroy_window()

def main():
    SAR_folder = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR"

    x_range = (-45, 45)
    y_range = (0, 60)
    
    os.makedirs(output_folder, exist_ok=True)

    with ThreadPoolExecutor() as executor:
        for root, _, files in os.walk(SAR_folder):
            files.sort(key=natural_key)
            for i in range(0, len(files), 1):
                file = files[i]
                if file.endswith(".mat"):
                    file_path = os.path.join(root, file)
                    print(f"正在处理文件: {file_path}")

                    mat_data = read_mat_file(file_path)
                    if mat_data:
                        executor.submit(
                            process_mat_file,
                            mat_data, 
                            x_range=x_range,
                            y_range=y_range
                        )

    plot_3d()

if __name__ == "__main__":
    main()